# Copyright 2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""TREC2014 Dataset"""
from pathlib import Path
import numpy as np
import argparse
import os
from scipy.sparse import csr_matrix
import random

from mindspore_gl.graph import MindHomoGraph, CsrAdj
from mindspore_gl.sampling.neighbor import sage_sampler_on_homo
from mindspore_gl import GraphField
import mindspore as ms
from src.base_dataset import BaseDataSet
from src.dataset import Dataset
# from trainval import args
# from base_dataset import BaseDataSet
# from dataset import Dataset


def qid_(dataset):
    
    qid_index = np.load('/disk1/gl/ms_new/src/data/dgat_qid_edge_index.npy')

    row_qid = qid_index[0][:]
    col_qid = qid_index[1][:]


    qid = np.ones(row_qid.shape)
    qid_node_count = dataset.query_size
    qid_csr_mat = csr_matrix((qid, (row_qid, col_qid)), shape=(qid_node_count+1, qid_node_count+1))
    qid_graph = MindHomoGraph()
    qid_node_dict = {idx: idx for idx in range(qid_node_count)}
    qid_edge_count = col_qid.shape[0]
    qid_edge_ids = np.array(list(range(qid_edge_count))).astype(np.int32)
    qid_graph.set_topo(CsrAdj(qid_csr_mat.indptr.astype(np.int32), qid_csr_mat.indices.astype(np.int32)),
                       qid_node_dict, qid_edge_ids)
    return qid_graph


def uid_(dataset):
    uid_index = np.load('/disk1/gl/ms_new/src/data/dgat_uid_edge_index.npy')

    row_uid = uid_index[0][:]
    col_uid = uid_index[1][:]
    uid = np.ones(row_uid.shape)

    uid_node_count = dataset.doc_size


    uid_csr_mat = csr_matrix((uid, (row_uid, col_uid)), shape=(uid_node_count+1, uid_node_count+1))
    uid_graph = MindHomoGraph()
    uid_node_dict = {idx: idx for idx in range(uid_node_count)}
    uid_edge_count = row_uid.shape[0]
    uid_edge_ids = np.array(list(range(uid_edge_count))).astype(np.int32)
    uid_graph.set_topo(CsrAdj(uid_csr_mat.indptr.astype(np.int32), uid_csr_mat.indices.astype(np.int32)),
                       uid_node_dict, uid_edge_ids)
    return uid_graph


def dgat_uid_neighor():
    with np.load('/disk1/gl/ms_new/src/data/data_neibor.npz', allow_pickle=True) as data:
        uid_neigh = data['uid_neigh']
        uid_num = data['uid_num']
        max_node_degree = data['max_node_degree']
    uid_neigh_sampler = ms.nn.Embedding(10960, 64)
    for idx, adj in enumerate(uid_neigh):
        adj_list = list(adj)
        if len(adj_list) >= max_node_degree:
            adj_sample = ms.Tensor.from_numpy(np.array(random.sample(adj_list, max_node_degree), dtype=np.int64))
        else:
            adj_sample = ms.Tensor.from_numpy(np.array(random.choices(adj_list, k=max_node_degree), dtype=np.int64))
        uid_neigh_sampler.embedding_table.data[idx] = adj_sample
    return uid_neigh_sampler


def qid_nei(dataset):
    qid_nodess = np.arange(0, dataset.query_size)
    qid_neighbor_sampler = sage_sampler_on_homo(homo_graph=qid_(dataset),
                                                seeds=qid_nodess.astype(np.int32),
                                                neighbor_nums=[5])

   
            
    qid_nodes = dataset.query_size
    qid_embedding = ms.nn.Embedding(qid_nodes, 64)
    qid_edges = qid_neighbor_sampler['layered_edges_0'].shape[1]
    qid_node = qid_neighbor_sampler['all_nodes']
    qid_node_embedding = qid_embedding(ms.Tensor(qid_node))
    qid_feat_size = 64  # 为embadding的后一位，将数据放入embadding其实shape的后一个维度大小为数据特征值
    qid_src_idx = ms.Tensor(qid_neighbor_sampler['layered_edges_0'][0][:])
    qid_dst_idx = ms.Tensor(qid_neighbor_sampler['layered_edges_0'][1][:])
    qid_ones = ms.ops.Ones()
    qid_feat = qid_ones((qid_nodes, qid_feat_size), ms.float32)
    qid_graph_field = GraphField(qid_src_idx, qid_dst_idx, qid_nodes, qid_edges)
    
    return qid_graph_field, qid_node_embedding, qid_node


def uid_nei(dataset):
    uid_nodes = dataset.doc_size
    uid_embedding = ms.nn.Embedding(uid_nodes, 64)
    uid_nodess = np.arange(0, dataset.doc_size)
    uid_neighbor_sampler = sage_sampler_on_homo(homo_graph=uid_(dataset),
                                                seeds=uid_nodess.astype(np.int32),
                                                neighbor_nums=[5])
    uid_edges = uid_neighbor_sampler['layered_edges_0'].shape[1]
    uid_node = uid_neighbor_sampler['all_nodes']
    uid_node_embedding = uid_embedding(ms.Tensor(uid_node))
    uid_feat_size = 64 #为embadding的后一位，将数据放入embadding其实shape的后一个维度大小为数据特征值
    uid_src_idx = ms.Tensor(uid_neighbor_sampler['layered_edges_0'][0][:])
    uid_dst_idx = ms.Tensor(uid_neighbor_sampler['layered_edges_0'][1][:])
    uid_ones = ms.ops.Ones()
    uid_feat = uid_ones((uid_nodes, uid_feat_size), ms.float32)
    uid_graph_field = GraphField(uid_src_idx, uid_dst_idx, uid_nodes, uid_edges)
    
    return uid_graph_field, uid_node_embedding, uid_node